import requests
from tipsz8 import saveLocalData, loadLocalData, clear_cache, getFileUTime
import time
import json
# Xinference 服务地址
url = "http://43.136.46.45:9997/v1/rerank"  # 替换为你的服务器地址
model_name = "bge-reranker-v2-m3"  # 模型名按你部署时的实际情况

# 输入 Query 和候选文本
# query = "推荐一款22000元左右的轻薄本"
# documents = [
#     "这是一款价格2399元的办公台式机",
#     "这款轻薄笔记本售价22999元，搭载13代酷睿",
#     "这台主机售价2599元，适合入门游戏"
# ]
def resort(query, documents):
    start_time = time.time()
    arrtemp = [json.dumps(i, ensure_ascii=False) for i in documents]
    # 构造请求体
    payload = {
        "model": model_name,
        "query": query,
        "documents": arrtemp
    }

    # 发起请求
    response = requests.post(url, json=payload)
    end_time = time.time()
    print(f"Response time: {end_time - start_time:.4f} seconds")
    # 解析结果
    if response.status_code == 200:
        result_json = response.json()
        reranked = result_json["results"]
        # print(reranked)
        # 将 index 映射回原文档
        sorted_results = sorted(reranked, key=lambda x: -x["relevance_score"])
        arr = []
        top_k = 5
        # ranked = sorted(zip(documents, sorted_results), key=lambda x: x[1], reverse=True)
        top_documents = sorted_results[:top_k]
        for rank, item in enumerate(top_documents):
            idx = item["index"]
            score = item["relevance_score"]
            text = documents[idx]
            # print(text["ID"])
            # arr.append(f"Rank {rank+1}: Score={score:.4f}, Text={text}")
            
            arr.append(text["ID"])
        # saveLocalData(arr, "rerank_result.json")
        return arr
        
  
        
            # print(f"Rank {rank+1}: Score={score:.4f}, Text={text}")
    else:
        print("Error:", response.status_code, response.text)